#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
AKShare请求工具模块
用于处理AKShare高频请求导致的连接问题
提供了请求限流、指数退避重试和缓存机制
"""

import time
import random
import logging
import functools
import threading
from typing import Any, Callable, Dict, Optional, Tuple, Union, List

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger('akshare_request')

# 全局请求限流器
class RequestThrottler:
    """请求限流器，控制请求频率"""
    
    def __init__(self, requests_per_minute: int = 120, burst: int = 5):
        """
        初始化请求限流器
        
        Args:
            requests_per_minute: 每分钟允许的请求数（默认降低到30次/分钟，避免连接被拒绝）
            burst: 允许的突发请求数（默认降低到5，避免短时间内发送过多请求）
        """
        self.rate = requests_per_minute / 60.0  # 每秒请求数
        self.burst = burst
        self.tokens = burst  # 初始令牌数
        self.last_update = time.time()
        self.lock = threading.RLock()
    
    def acquire(self, block: bool = True) -> bool:
        """
        获取请求令牌
        
        Args:
            block: 是否阻塞等待令牌
            
        Returns:
            bool: 是否获取到令牌
        """
        with self.lock:
            now = time.time()
            # 计算新增令牌
            new_tokens = (now - self.last_update) * self.rate
            self.tokens = min(self.burst, self.tokens + new_tokens)
            self.last_update = now
            
            if self.tokens >= 1:
                # 有可用令牌
                self.tokens -= 1
                return True
            elif not block:
                # 无令牌且不阻塞
                return False
            
        # 需要等待
        wait_time = (1.0 - self.tokens) / self.rate
        time.sleep(wait_time)
        return self.acquire(block=False)

# 创建全局限流器实例
_global_throttler = RequestThrottler()

def retry_with_exponential_backoff(max_retries: int = 2, 
                                  base_delay: float = 1.0,
                                  max_delay: float = 3.0,
                                  jitter: bool = True,
                                  error_types: Tuple = (Exception,)):
    """
    指数退避重试装饰器
    
    Args:
        max_retries: 最大重试次数
        base_delay: 基础延迟时间（秒）
        max_delay: 最大延迟时间（秒）
        jitter: 是否添加随机抖动
        error_types: 需要重试的异常类型
        
    Returns:
        装饰器函数
    """
    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            retries = 0
            while True:
                try:
                    # 获取请求令牌（限流）
                    _global_throttler.acquire()
                    
                    # 执行请求
                    return func(*args, **kwargs)
                    
                except error_types as e:
                    retries += 1
                    if retries > max_retries:
                        logger.error(f"达到最大重试次数 {max_retries}，放弃重试: {str(e)}")
                        raise
                    
                    # 计算延迟时间
                    delay = min(base_delay * (2 ** (retries - 1)), max_delay)
                    if jitter:
                        delay = delay * (0.5 + random.random())
                    
                    logger.warning(f"请求失败 (尝试 {retries}/{max_retries}): {str(e)[:100]}... 等待 {delay:.2f} 秒后重试")
                    time.sleep(delay)
        return wrapper
    return decorator

# 创建内存缓存
_cache = {}
_cache_lock = threading.RLock()

def cached(ttl: int = 3600):
    """
    缓存装饰器
    
    Args:
        ttl: 缓存生存时间（秒）
        
    Returns:
        装饰器函数
    """
    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            # 生成缓存键
            cache_key = f"{func.__name__}:{str(args)}:{str(kwargs)}"
            
            with _cache_lock:
                # 检查缓存
                if cache_key in _cache:
                    timestamp, value = _cache[cache_key]
                    if time.time() - timestamp < ttl:
                        return value
            
            # 执行函数
            result = func(*args, **kwargs)
            
            # 更新缓存
            with _cache_lock:
                _cache[cache_key] = (time.time(), result)
                
                # 清理过期缓存
                now = time.time()
                expired_keys = [k for k, (ts, _) in _cache.items() if now - ts > ttl]
                for k in expired_keys:
                    del _cache[k]
            
            return result
        return wrapper
    return decorator

def safe_akshare_request(func: Callable, *args, **kwargs) -> Any:
    """
    安全执行AKShare请求的通用函数
    
    Args:
        func: AKShare函数
        *args: 位置参数
        **kwargs: 关键字参数
        
    Returns:
        Any: 请求结果
    """
    @retry_with_exponential_backoff()
    def _execute():
        return func(*args, **kwargs)
    
    return _execute()

# 为常用AKShare函数提供包装器
def get_stock_zh_a_hist_safe(symbol: str, period: str = "daily", 
                           start_date: str = None, end_date: str = None, 
                           adjust: str = "qfq", **kwargs) -> Any:
    """
    安全获取A股历史数据
    
    Args:
        symbol: 股票代码
        period: 周期
        start_date: 开始日期
        end_date: 结束日期
        adjust: 复权类型
        **kwargs: 其他参数
        
    Returns:
        Any: 股票历史数据
    """
    import akshare as ak
    
    # 生成唯一的缓存键，包含所有参数
    cache_key = f"stock_hist:{symbol}:{period}:{start_date}:{end_date}:{adjust}:{str(kwargs)}"
    
    # 检查缓存
    with _cache_lock:
        if cache_key in _cache:
            timestamp, value = _cache[cache_key]
            if time.time() - timestamp < 3600:  # 1小时缓存
                return value
    
    # 如果缓存未命中，执行请求
    @retry_with_exponential_backoff(max_retries=2, base_delay=1, max_delay=2)
    def _get_stock_history():
        # 获取请求令牌（限流）
        _global_throttler.acquire()
        
        # 添加预请求等待，减少连接被拒绝的可能性
        time.sleep(random.uniform(1.0, 3.0))
        
        # 执行请求
        return ak.stock_zh_a_hist(symbol=symbol, period=period, 
                                 start_date=start_date, end_date=end_date, 
                                 adjust=adjust, **kwargs)
    
    # 获取结果
    result = _get_stock_history()
    
    # 更新缓存
    with _cache_lock:
        _cache[cache_key] = (time.time(), result)
    
    return result

@cached(ttl=3600 * 24) # 缓存24小时
def get_stock_zh_a_hist_safe1(symbol: str, period: str = "daily", 
                           start_date: str = None, end_date: str = None, 
                           adjust: str = "qfq", **kwargs) -> Any:
    """
    安全获取A股历史数据
    
    Args:
        symbol: 股票代码
        period: 周期
        start_date: 开始日期
        end_date: 结束日期
        adjust: 复权类型
        **kwargs: 其他参数
        
    Returns:
        Any: 股票历史数据
    """
    import akshare as ak # 确保ak在函数作用域内可用
    logger.info(f"正在从AKShare获取股票 {symbol} 的历史数据...")
    try:
        df = safe_akshare_request(
            func=ak.stock_zh_a_hist,
            symbol=symbol,
            period=period,
            start_date=start_date,
            end_date=end_date,
            adjust=adjust,
            **kwargs
        )
        if df is None or df.empty:
            logger.warning(f"股票 {symbol} 在 {start_date} 到 {end_date} 期间没有获取到数据。")
            return None
        return df
    except Exception as e:
        logger.error(f"获取股票 {symbol} 历史数据失败: {e}")
        return None
    
def get_stock_zh_a_spot_em_safe(**kwargs) -> Any:
    """
    安全获取A股实时行情
    
    Args:
        **kwargs: 其他参数
        
    Returns:
        Any: A股实时行情数据
    """
    import akshare as ak
    
    # 生成唯一的缓存键，包含所有参数
    cache_key = f"stock_spot:{str(kwargs)}"
    
    # 检查缓存
    with _cache_lock:
        if cache_key in _cache:
            timestamp, value = _cache[cache_key]
            if time.time() - timestamp < 300:  # 5分钟缓存
                return value
    
    # 如果缓存未命中，执行请求
    #  @retry_with_exponential_backoff(max_retries=10, base_delay=2.0, max_delay=120.0)
    @retry_with_exponential_backoff(max_retries=2, base_delay=1, max_delay=2)
    def _get_stock_spot():
        # 获取请求令牌（限流）
        _global_throttler.acquire()
        
        # 添加预请求等待，减少连接被拒绝的可能性
        time.sleep(random.uniform(1.0, 3.0))
        
        # 执行请求
        return ak.stock_zh_a_spot_em(**kwargs)
    
    # 获取结果
    result = _get_stock_spot()
    
    # 更新缓存
    with _cache_lock:
        _cache[cache_key] = (time.time(), result)
    
    return result

def get_tool_trade_date_hist_sina_safe(**kwargs) -> Any:
    """
    安全获取新浪交易日历
    
    Args:
        **kwargs: 其他参数
        
    Returns:
        Any: 交易日历数据
    """
    import akshare as ak
    
    # 生成唯一的缓存键，包含所有参数
    cache_key = f"trade_date:{str(kwargs)}"
    
    # 检查缓存
    with _cache_lock:
        if cache_key in _cache:
            timestamp, value = _cache[cache_key]
            if time.time() - timestamp < 86400:  # 24小时缓存
                return value
    
    # 如果缓存未命中，执行请求
    @retry_with_exponential_backoff(max_retries=2)
    def _get_trade_date():
        # 获取请求令牌（限流）
        _global_throttler.acquire()
        
        # 添加预请求等待，减少连接被拒绝的可能性
        time.sleep(random.uniform(1.0, 3.0))
        
        # 执行请求
        return ak.tool_trade_date_hist_sina(**kwargs)
    
    # 获取结果
    result = _get_trade_date()
    
    # 更新缓存
    with _cache_lock:
        _cache[cache_key] = (time.time(), result)
    
    return result

# 示例用法
if __name__ == "__main__":
    import akshare as ak
    import pandas as pd
    
    # 测试安全请求函数
    try:
        # 获取交易日历
        trade_cal = get_tool_trade_date_hist_sina_safe()
        print(f"获取交易日历成功，共 {len(trade_cal)} 条记录")
        
        # 获取股票列表
        stock_list = get_stock_zh_a_spot_em_safe()
        print(f"获取股票列表成功，共 {len(stock_list)} 只股票")
        
        # 获取单只股票历史数据
        if not stock_list.empty:
            stock_code = stock_list.iloc[0]["代码"]
            history = get_stock_zh_a_hist_safe(symbol=stock_code, start_date="20230101", end_date="20230110")
            print(f"获取股票 {stock_code} 历史数据成功，共 {len(history)} 条记录")
    except Exception as e:
        print(f"测试失败: {e}")